//
// Created by yangchao on 2022/6/1.
// 105. 从前序与中序遍历序列构造二叉树: https://leetcode.cn/problems/construct-binary-tree-from-preorder-and-inorder-traversal/
//

#include "../DataStructure/TreeNode.h"
#include <vector>
#include <stack>
#include <unordered_map>
using namespace std;

class BuildTree {
public:
    //1. 递归
    TreeNode* buildTree(vector<int>& preorder, vector<int>& inorder) {
        unordered_map <int,int> mp;
        int len = inorder.size();
        for (int i = 0; i < len; ++i) {
            //键表示一个元素（节点的值），值表示其在中序遍历中的出现位置
            mp.insert(pair<int,int>(inorder[i],i));
            //mp[inorder[i]] = i;
        }
        return buildTreeHelper(preorder,0,len-1,inorder,0,len-1,mp);
    }

    //2. 迭代
    TreeNode* buildTreeII(vector<int>& preorder, vector<int>& inorder) {
        if (!preorder.size()) return nullptr;
        TreeNode* root = new TreeNode(preorder[0]);
        stack<TreeNode*> stk;
        stk.push(root);
        int inorderIndex = 0;
        for (int i = 1; i < preorder.size(); ++i) {
            int preorderVal = preorder[i];
            TreeNode* node = stk.top();
            if (node->val != inorder[inorderIndex]) {
                node->left = new TreeNode(preorderVal);
                stk.push(node->left);
            } else{
                while (!stk.empty() && stk.top()->val == inorder[inorderIndex]) {
                    node = stk.top();
                    stk.pop();
                    ++inorderIndex;
                }
                node->right = new TreeNode(preorderVal);
                stk.push(node->right);
            }
        }
        return root;
    }

private:
    TreeNode* buildTreeHelper(vector<int>& preorder, int p_start, int p_end,
                              vector<int>& inorder, int i_start, int i_end, unordered_map<int,int>& mp) {
        if (p_start > p_end) return nullptr;
        //前序遍历的第一个节点就是根节点
        int rootVal = preorder[p_start];
        TreeNode *root = new TreeNode(rootVal);
        //在中序遍历中定位根节点
        int i_root_index = mp[rootVal];
        //得到左子树中节点的数目
        int leftNum = i_root_index - i_start;
        root->left = buildTreeHelper(preorder, p_start + 1, p_start + leftNum, inorder, i_start, i_root_index - 1, mp);
        root->right = buildTreeHelper(preorder, p_start + leftNum + 1, p_end, inorder, i_root_index + 1, i_end, mp);
        return root;
    }
};